from fastapi import FastAPI
from typing import Callable
from config import settings
from services.storage import LocalFileStorage, AiocacheMemory, BaseStorage, MinIoObjectStorage
from services import AiModelCache
import logging
import argparse
logger = logging.getLogger(__name__)


async def _startup_model(app: FastAPI, args:argparse.Namespace) -> None:
    model_path = settings.MODEL_PATH
    logger.info("HF Model path:%s" % model_path)
    # 是否是本地
    app.state.storage = init_storage()
    #
    amc = AiModelCache(AiocacheMemory())
    # 加载的模型别名配置
    loadModelConfig:dict = {'nlp': args.nlp, 'tts': args.tts, 'asr': args.asr}
    await amc.load_model(device='cuda', **loadModelConfig)
    app.state.cache = amc

def init_storage()->BaseStorage:
    if settings.STORAGE == 'file':
        return LocalFileStorage(base_url=f'{settings.BASE_URL}/oss')
    elif settings.STORAGE == 'minio':
        return MinIoObjectStorage(endpoint=settings.MINIO_Endpoint, 
                                  access_key=settings.MINIO_Access_Key,
                                  secret_key=settings.MINIO_Secret_Key,
                                  is_https=settings.MINIO_Secure,
                                  bucket=settings.MINIO_Bucket)
    else:
        return None
    
def _shutdown_model(app: FastAPI) -> None:
    app.state.storage = None
    app.state.cache = None


def start_app_handler(app: FastAPI, args:argparse.Namespace) -> Callable:
    async def startup() -> None:
        logger.info("Running app start handler.")
        await _startup_model(app, args)

    return startup


def stop_app_handler(app: FastAPI) -> Callable:
    def shutdown() -> None:
        logger.info("Running app shutdown handler.")
        _shutdown_model(app)

    return shutdown
